// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <cmath>
#include <queue>
#include <utility>

#include <nebula/compute/api_aggregate.h>
#include <nebula/compute/kernels/aggregate_internal.h>
#include <nebula/compute/kernels/common_internal.h>
#include <nebula/compute/kernels/util_internal.h>

#include <nebula/core/stl_allocator.h>
#include <nebula/types/type_traits.h>
#include <nebula/bits/bit_util.h>

namespace nebula::compute::internal {

    namespace {

        using ModeState = OptionsWrapper<ModeOptions>;

        constexpr char kModeFieldName[] = "mode";
        constexpr char kCountFieldName[] = "count";

        constexpr uint64_t kCountEOF = ~0ULL;

        template<typename InType, typename CType = typename TypeTraits<InType>::CType>
        turbo::Result<std::pair<CType *, int64_t *>> PrepareOutput(int64_t n, KernelContext *ctx,
                                                                   const DataType &type, ExecResult *out) {
            DKCHECK_EQ(Type::STRUCT, type.id());
            const auto &out_type = turbo::checked_cast<const StructType &>(type);
            DKCHECK_EQ(2, out_type.num_fields());
            const auto &mode_type = out_type.field(0)->type();
            const auto &count_type = int64();

            auto mode_data = ArrayData::create(mode_type, /*length=*/n, /*null_count=*/0);
            mode_data->buffers.resize(2, nullptr);
            auto count_data = ArrayData::create(count_type, n, 0);
            count_data->buffers.resize(2, nullptr);

            CType *mode_buffer = nullptr;
            int64_t *count_buffer = nullptr;

            if (n > 0) {
                const auto mode_buffer_size = bit_util::BytesForBits(n * mode_type->bit_width());
                TURBO_MOVE_OR_RAISE(mode_data->buffers[1], ctx->allocate(mode_buffer_size));
                TURBO_MOVE_OR_RAISE(count_data->buffers[1], ctx->allocate(n * sizeof(int64_t)));
                mode_buffer = mode_data->template GetMutableValues<CType>(1);
                count_buffer = count_data->template GetMutableValues<int64_t>(1);
            }

            out->value =
                    ArrayData::create(type.get_shared_ptr(), n, {nullptr}, {mode_data, count_data}, 0);
            return std::make_pair(mode_buffer, count_buffer);
        }

// find top-n value:count pairs with minimal heap
// suboptimal for tiny or large n, possibly okay as we're not in hot path
        template<typename InType, typename Generator>
        turbo::Status Finalize(KernelContext *ctx, const DataType &type, ExecResult *out,
                               Generator &&gen) {
            using CType = typename TypeTraits<InType>::CType;

            using ValueCountPair = std::pair<CType, uint64_t>;
            auto gt = [](const ValueCountPair &lhs, const ValueCountPair &rhs) {
                const bool rhs_is_nan = rhs.first != rhs.first;  // nan as largest value
                return lhs.second > rhs.second ||
                       (lhs.second == rhs.second && (lhs.first < rhs.first || rhs_is_nan));
            };

            std::priority_queue<ValueCountPair, std::vector<ValueCountPair>, decltype(gt)> min_heap(
                    std::move(gt));

            const ModeOptions &options = ModeState::Get(ctx);
            while (true) {
                const ValueCountPair &value_count = gen();
                DKCHECK_NE(value_count.second, 0);
                if (value_count.second == kCountEOF) break;
                if (static_cast<int64_t>(min_heap.size()) < options.n) {
                    min_heap.push(value_count);
                } else if (gt(value_count, min_heap.top())) {
                    min_heap.pop();
                    min_heap.push(value_count);
                }
            }
            const int64_t n = min_heap.size();

            CType *mode_buffer;
            int64_t *count_buffer;
            TURBO_MOVE_OR_RAISE(std::tie(mode_buffer, count_buffer),
                                PrepareOutput<InType>(n, ctx, type, out));

            for (int64_t i = n - 1; i >= 0; --i) {
                std::tie(mode_buffer[i], count_buffer[i]) = min_heap.top();
                min_heap.pop();
            }

            return turbo::OkStatus();
        }

// count value occurrences for integers with narrow value range
// O(1) space, O(n) time
        template<typename T>
        struct CountModer {
            using CType = typename T::c_type;

            CType min;
            std::vector<uint64_t> counts;

            CountModer(CType min, CType max) {
                uint32_t value_range = static_cast<uint32_t>(max - min) + 1;
                DKCHECK_LT(value_range, 1 << 20);
                this->min = min;
                this->counts.resize(value_range, 0);
            }

            turbo::Status get_result(KernelContext *ctx, const DataType &type, ExecResult *out) {
                // generator to emit next value:count pair
                int index = 0;
                auto gen = [&]() {
                    for (; index < static_cast<int>(counts.size()); ++index) {
                        if (counts[index] != 0) {
                            auto value_count =
                                    std::make_pair(static_cast<CType>(index + this->min), counts[index]);
                            ++index;
                            return value_count;
                        }
                    }
                    return std::pair<CType, uint64_t>(0, kCountEOF);
                };

                return Finalize<T>(ctx, type, out, std::move(gen));
            }

            turbo::Status Exec(KernelContext *ctx, const ExecSpan &batch, ExecResult *out) {
                // count values in all chunks, ignore nulls
                const ArraySpan &values = batch[0].array;
                const ModeOptions &options = ModeState::Get(ctx);
                if ((!options.skip_nulls && values.get_null_count() > 0) ||
                    (values.length - values.get_null_count() < options.min_count)) {
                    return PrepareOutput<T>(/*n=*/0, ctx, *out->type(), out).status();
                }

                CountValues<CType>(values, this->min, this->counts.data());
                return get_result(ctx, *out->type(), out);
            }

            turbo::Status exec_chunked(KernelContext *ctx, const ExecBatch &batch, Datum *out) {
                // count values in all chunks, ignore nulls
                const ChunkedArray &values = *batch[0].chunked_array();
                const ModeOptions &options = ModeState::Get(ctx);
                ExecResult result;
                if ((!options.skip_nulls && values.null_count() > 0) ||
                    (values.length() - values.null_count() < options.min_count)) {
                    TURBO_RETURN_NOT_OK(PrepareOutput<T>(/*n=*/0, ctx, *out->type(), &result));
                } else {
                    CountValues<CType>(values, this->min, this->counts.data());
                    TURBO_RETURN_NOT_OK(get_result(ctx, *out->type(), &result));
                }
                *out = result.array_data();
                return turbo::OkStatus();
            }
        };

// booleans can be handled more straightforward
        template<>
        struct CountModer<BooleanType> {
            int64_t counts[2] = {0, 0};

            void UpdateCounts(const ArraySpan &values) {
                if (values.length > values.get_null_count()) {
                    const int64_t true_count = GetTrueCount(values);
                    counts[true] += true_count;
                    counts[false] += values.length - values.null_count - true_count;
                }
            }

            void UpdateCounts(const ChunkedArray &values) {
                for (const auto &chunk: values.chunks()) {
                    UpdateCounts(*chunk->data());
                }
            }

            turbo::Status WrapResult(KernelContext *ctx, const ModeOptions &options, const DataType &type,
                                     ExecResult *out) {
                const int64_t distinct_values = (this->counts[0] != 0) + (this->counts[1] != 0);
                const int64_t n = std::min(options.n, distinct_values);

                uint8_t *mode_buffer;
                int64_t *count_buffer;
                TURBO_MOVE_OR_RAISE(std::tie(mode_buffer, count_buffer),
                                    (PrepareOutput<BooleanType, uint8_t>(n, ctx, type, out)));

                if (n >= 1) {
                    // at most two bits are useful in mode buffer
                    mode_buffer[0] = 0;
                    const bool first_mode = counts[true] > counts[false];
                    bit_util::SetBitTo(mode_buffer, 0, first_mode);
                    count_buffer[0] = counts[first_mode];
                    if (n == 2) {
                        const bool second_mode = !first_mode;
                        bit_util::SetBitTo(mode_buffer, 1, second_mode);
                        count_buffer[1] = counts[second_mode];
                    }
                }

                return turbo::OkStatus();
            }

            turbo::Status Exec(KernelContext *ctx, const ExecSpan &batch, ExecResult *out) {
                const ArraySpan &values = batch[0].array;
                const ModeOptions &options = ModeState::Get(ctx);
                if ((!options.skip_nulls && values.get_null_count() > 0) ||
                    (values.length - values.null_count < options.min_count)) {
                    return PrepareOutput<BooleanType, uint8_t>(0, ctx, *out->type(), out).status();
                }
                UpdateCounts(values);
                return WrapResult(ctx, options, *out->type(), out);
            }

            turbo::Status exec_chunked(KernelContext *ctx, const ExecBatch &batch, Datum *out) {
                const ChunkedArray &values = *batch[0].chunked_array();
                const ModeOptions &options = ModeState::Get(ctx);
                ExecResult result;
                if ((!options.skip_nulls && values.null_count() > 0) ||
                    (values.length() - values.null_count() < options.min_count)) {
                    TURBO_RETURN_NOT_OK((PrepareOutput<BooleanType, uint8_t>(0, ctx, *out->type(), &result)));
                } else {
                    UpdateCounts(values);
                    TURBO_RETURN_NOT_OK(WrapResult(ctx, options, *out->type(), &result));
                }
                *out = result.array_data();
                return turbo::OkStatus();
            }
        };

// copy and sort approach for floating points, decimals, or integers with wide
// value range
// O(n) space, O(nlogn) time
        template<typename T>
        struct SortModer {
            using CType = typename TypeTraits<T>::CType;
            using Allocator = nebula::stl::allocator<CType>;

            template<typename Type = T>
            static enable_if_floating_point<Type, CType> GetNan() {
                return static_cast<CType>(NAN);
            }

            template<typename Type = T>
            static enable_if_t<!is_floating_type<Type>::value, CType> GetNan() {
                DKCHECK(false);
                return static_cast<CType>(0);
            }

            template<typename Container>
            turbo::Status ComputeMode(KernelContext *ctx, const Container &arr, int64_t length,
                                      int64_t null_count, const DataType &type, ExecResult *out) {
                const ModeOptions &options = ModeState::Get(ctx);
                const int64_t in_length = length - null_count;
                if ((!options.skip_nulls && null_count > 0) || (in_length < options.min_count)) {
                    return PrepareOutput<T>(/*n=*/0, ctx, type, out).status();
                }

                // copy all chunks to a buffer, ignore nulls and nans
                std::vector<CType, Allocator> values(Allocator(ctx->memory_pool()));

                uint64_t nan_count = 0;
                if (length > 0) {
                    values.resize(length - null_count);
                    CopyNonNullValues(arr, values.data());

                    // drop nan
                    if (is_floating_type<T>::value) {
                        const auto &it =
                                std::remove_if(values.begin(), values.end(), [](CType v) { return v != v; });
                        nan_count = values.end() - it;
                        values.resize(it - values.begin());
                    }
                }
                // sort the input data to count same values
                std::sort(values.begin(), values.end());

                // generator to emit next value:count pair
                auto it = values.cbegin();
                auto gen = [&]() {
                    if (TURBO_UNLIKELY(it == values.cend())) {
                        // handle NAN at last
                        if (nan_count > 0) {
                            auto value_count = std::make_pair(GetNan(), nan_count);
                            nan_count = 0;
                            return value_count;
                        }
                        return std::pair<CType, uint64_t>(static_cast<CType>(0), kCountEOF);
                    }
                    // count same values
                    const CType value = *it;
                    uint64_t count = 0;
                    do {
                        ++it;
                        ++count;
                    } while (it != values.cend() && *it == value);
                    return std::make_pair(value, count);
                };

                return Finalize<T>(ctx, type, out, std::move(gen));
            }

            turbo::Status Exec(KernelContext *ctx, const ExecSpan &batch, ExecResult *out) {
                const ArraySpan &values = batch[0].array;
                return ComputeMode(ctx, values, values.length, values.get_null_count(), *out->type(),
                                   out);
            }

            turbo::Status exec_chunked(KernelContext *ctx, const ExecBatch &batch, Datum *out) {
                const ChunkedArray &values = *batch[0].chunked_array();
                ExecResult result;
                TURBO_RETURN_NOT_OK(ComputeMode(ctx, values, values.length(), values.null_count(),
                                                *out->type(), &result));
                *out = result.array_data();
                return turbo::OkStatus();
            }
        };

        template<typename CType, typename Container>
        bool ShouldUseCountMode(const Container &values, int64_t num_valid, CType *min,
                                CType *max) {
            // cross point to benefit from counting approach
            // about 2x improvement for int32/64 from micro-benchmarking
            static constexpr int kMinArraySize = 8192;
            static constexpr int kMaxValueRange = 32768;

            if (num_valid >= kMinArraySize) {
                std::tie(*min, *max) = GetMinMax<CType>(values);
                return static_cast<uint64_t>(*max) - static_cast<uint64_t>(*min) <= kMaxValueRange;
            }
            return false;
        }

// pick counting or sorting approach per integers value range
        template<typename T>
        struct CountOrSortModer {
            using CType = typename T::c_type;

            turbo::Status Exec(KernelContext *ctx, const ExecSpan &batch, ExecResult *out) {
                const ArraySpan &values = batch[0].array;
                CType min, max;
                if (ShouldUseCountMode<CType>(values, values.length - values.get_null_count(), &min,
                                              &max)) {
                    return CountModer<T>(min, max).Exec(ctx, batch, out);
                }
                return SortModer<T>().Exec(ctx, batch, out);
            }

            turbo::Status exec_chunked(KernelContext *ctx, const ExecBatch &batch, Datum *out) {
                const ChunkedArray &values = *batch[0].chunked_array();
                CType min, max;
                if (ShouldUseCountMode<CType>(values, values.length() - values.null_count(), &min,
                                              &max)) {
                    return CountModer<T>(min, max).exec_chunked(ctx, batch, out);
                }
                return SortModer<T>().exec_chunked(ctx, batch, out);
            }
        };

        template<typename InType, typename Enable = void>
        struct Moder;

        template<>
        struct Moder<Int8Type> {
            CountModer<Int8Type> impl;

            Moder() : impl(-128, 127) {}
        };

        template<>
        struct Moder<UInt8Type> {
            CountModer<UInt8Type> impl;

            Moder() : impl(0, 255) {}
        };

        template<>
        struct Moder<BooleanType> {
            CountModer<BooleanType> impl;
        };

        template<typename InType>
        struct Moder<InType, enable_if_t<(is_integer_type<InType>::value &&
                                          (sizeof(typename InType::c_type) > 1))>> {
            CountOrSortModer<InType> impl;
        };

        template<typename InType>
        struct Moder<InType, enable_if_floating_point<InType>> {
            SortModer<InType> impl;
        };

        template<typename InType>
        struct Moder<InType, enable_if_decimal<InType>> {
            SortModer<InType> impl;
        };

        turbo::Status CheckOptions(KernelContext *ctx) {
            if (ctx->state() == nullptr) {
                return turbo::invalid_argument_error("Mode requires ModeOptions");
            }
            const ModeOptions &options = ModeState::Get(ctx);
            if (options.n <= 0) {
                return turbo::invalid_argument_error("ModeOptions::n must be strictly positive");
            }
            return turbo::OkStatus();
        }

        template<typename OutTypeUnused, typename InType>
        struct ModeExecutor {
            static turbo::Status Exec(KernelContext *ctx, const ExecSpan &batch, ExecResult *out) {
                TURBO_RETURN_NOT_OK(CheckOptions(ctx));
                return Moder<InType>().impl.Exec(ctx, batch, out);
            }
        };

        template<typename OutTypeUnused, typename InType>
        struct ModeExecutorChunked {
            static turbo::Status Exec(KernelContext *ctx, const ExecBatch &batch, Datum *out) {
                TURBO_RETURN_NOT_OK(CheckOptions(ctx));
                return Moder<InType>().impl.exec_chunked(ctx, batch, out);
            }
        };

        turbo::Result<TypeHolder> ModeType(KernelContext *, const std::vector<TypeHolder> &types) {
            return STRUCT(
                    {field(kModeFieldName, types[0].get_shared_ptr()), field(kCountFieldName, int64())});
        }

        VectorKernel NewModeKernel(const std::shared_ptr<DataType> &in_type, ArrayKernelExec exec,
                                   VectorKernel::ChunkedExec exec_chunked) {
            VectorKernel kernel;
            kernel.init = ModeState::init;
            kernel.can_execute_chunkwise = false;
            kernel.output_chunked = false;
            switch (in_type->id()) {
                case Type::DECIMAL128:
                case Type::DECIMAL256:
                    kernel.signature =
                            KernelSignature::create({InputType(in_type->id())}, OutputType(ModeType));
                    break;
                default: {
                    auto out_type =
                            STRUCT({field(kModeFieldName, in_type), field(kCountFieldName, int64())});
                    kernel.signature = KernelSignature::create({in_type->id()}, std::move(out_type));
                    break;
                }
            }
            kernel.exec = std::move(exec);
            kernel.exec_chunked = exec_chunked;
            return kernel;
        }

        const FunctionDoc mode_doc{
                "Compute the modal (most common) values of a numeric array",
                ("Compute the n most common values and their respective occurrence counts.\n"
                 "The output has type `struct<mode: T, count: int64>`, where T is the\n"
                 "input type.\n"
                 "The results are ordered by descending `count` first, and ascending `mode`\n"
                 "when breaking ties.\n"
                 "Nulls are ignored.  If there are no non-null values in the array,\n"
                 "an empty array is returned."),
                {"array"},
                "ModeOptions"};

    }  // namespace

    void register_scalar_aggregate_mode(FunctionRegistry *registry) {
        static auto default_options = ModeOptions::defaults();
        auto func = std::make_shared<VectorFunction>("mode", Arity::Unary(), mode_doc,
                                                     &default_options);
        KCHECK_OK(func->add_kernel(
                NewModeKernel(boolean(), ModeExecutor<StructType, BooleanType>::Exec,
                              ModeExecutorChunked<StructType, BooleanType>::Exec)));
        for (const auto &type: numeric_types()) {
            // TODO(wesm):
            KCHECK_OK(func->add_kernel(NewModeKernel(
                    type, generate_numeric<ModeExecutor, StructType>(*type),
                    generate_numeric<ModeExecutorChunked, StructType, VectorKernel::ChunkedExec>(
                            *type))));
        }
        // Type parameters are ignored
        KCHECK_OK(func->add_kernel(
                NewModeKernel(decimal128(1, 0), ModeExecutor<StructType, Decimal128Type>::Exec,
                              ModeExecutorChunked<StructType, Decimal128Type>::Exec)));
        KCHECK_OK(func->add_kernel(
                NewModeKernel(decimal256(1, 0), ModeExecutor<StructType, Decimal256Type>::Exec,
                              ModeExecutorChunked<StructType, Decimal256Type>::Exec)));
        KCHECK_OK(registry->add_function(std::move(func)));
    }

}  // namespace nebula::compute::internal
